// pc_multi_merger.cpp — Merge two PointCloud2 streams using adjustable extrinsics (C->A)
// - 保留所有字段，只对 C 的 x/y/z 做刚体变换，再与 A 逐字节拼接
// - 加入 dynamic_reconfigure 滑块 (tx,ty,tz,yaw,pitch,roll,apply_tf)
// - 可选实时发布 TF: child=frame_C, parent=frame_A

#include <ros/ros.h>
#include <sensor_msgs/PointCloud2.h>
#include <message_filters/subscriber.h>
#include <message_filters/synchronizer.h>
#include <message_filters/sync_policies/approximate_time.h>

#include <Eigen/Core>
#include <Eigen/Geometry>
#include <algorithm>
#include <string>
#include <vector>
#include <stdexcept>
#include <mutex>

// dynamic_reconfigure
#include <dynamic_reconfigure/server.h>
#include <pc_udp_streamer/ExtrinsicConfig.h>

// tf broadcaster
#include <tf2_ros/transform_broadcaster.h>
#include <geometry_msgs/TransformStamped.h>

struct FieldOffsets {
  int off_x{-1}, off_y{-1}, off_z{-1};
  uint32_t point_step{0};
  bool little_endian{true};
};

static inline double deg2rad(double d){ return d * M_PI / 180.0; }

static bool findXYZOffsets(const sensor_msgs::PointCloud2& msg, FieldOffsets& fo){
  fo.point_step = msg.point_step;
  fo.little_endian = (msg.is_bigendian == 0);
  for(const auto& f : msg.fields){
    if(f.name == "x") fo.off_x = f.offset;
    else if(f.name == "y") fo.off_y = f.offset;
    else if(f.name == "z") fo.off_z = f.offset;
  }
  return (fo.off_x>=0 && fo.off_y>=0 && fo.off_z>=0);
}

static void transformXYZ_in_place(sensor_msgs::PointCloud2& cloud,
                                  const FieldOffsets& fo,
                                  const Eigen::Matrix4f& T)
{
  const size_t n = static_cast<size_t>(cloud.width) * cloud.height;
  uint8_t* base = cloud.data.data();

  for(size_t i=0;i<n;++i){
    uint8_t* p = base + i * fo.point_step;
    float* px = reinterpret_cast<float*>(p + fo.off_x);
    float* py = reinterpret_cast<float*>(p + fo.off_y);
    float* pz = reinterpret_cast<float*>(p + fo.off_z);
    const Eigen::Vector4f v(*px, *py, *pz, 1.f);
    Eigen::Vector4f vt = T * v;
    *px = vt.x(); *py = vt.y(); *pz = vt.z();
  }
}

static sensor_msgs::PointCloud2 concatCloudsSameLayout(const sensor_msgs::PointCloud2& a,
                                                       const sensor_msgs::PointCloud2& b)
{
  if(a.fields.size()!=b.fields.size() || a.point_step!=b.point_step ||
     a.is_bigendian != b.is_bigendian){
    throw std::runtime_error("PointCloud2 layout mismatch (fields/step/endian)");
  }
  const uint32_t na = a.width * a.height;
  const uint32_t nb = b.width * b.height;

  sensor_msgs::PointCloud2 out;
  out.header = a.header;          // 输出沿用 A 的 header（frame 已是 A）
  out.height = 1;
  out.width  = na + nb;
  out.fields = a.fields;
  out.is_bigendian = a.is_bigendian;
  out.point_step = a.point_step;
  out.row_step = out.point_step * out.width;
  out.is_dense = a.is_dense && b.is_dense;
  out.data.resize(static_cast<size_t>(out.row_step));

  // 拷贝 A
  const size_t bytes_a = static_cast<size_t>(na) * a.point_step;
  if (na==a.width*a.height && a.row_step==a.point_step*a.width){
    std::memcpy(out.data.data(), a.data.data(), bytes_a);
  }else{
    size_t off=0;
    for(uint32_t r=0;r<a.height;++r){
      const uint8_t* row = a.data.data() + r * a.row_step;
      std::memcpy(out.data.data()+off, row, a.width * a.point_step);
      off += a.width * a.point_step;
    }
  }

  // 叠加 B
  const size_t bytes_b = static_cast<size_t>(nb) * b.point_step;
  if (nb==b.width*b.height && b.row_step==b.point_step*b.width){
    std::memcpy(out.data.data() + bytes_a, b.data.data(), bytes_b);
  }else{
    size_t off=bytes_a;
    for(uint32_t r=0;r<b.height;++r){
      const uint8_t* row = b.data.data() + r * b.row_step;
      std::memcpy(out.data.data()+off, row, b.width * b.point_step);
      off += b.width * b.point_step;
    }
  }
  return out;
}

class PcMultiMerger {
  using SyncPolicy = message_filters::sync_policies::ApproximateTime<
      sensor_msgs::PointCloud2, sensor_msgs::PointCloud2>;

public:
  explicit PcMultiMerger(ros::NodeHandle& nh, ros::NodeHandle& pnh)
  : sub_A_(nh, pnh.param<std::string>("in_cloud_A", "/rx_A"), pnh.param<int>("queue_size", 10)),
    sub_C_(nh, pnh.param<std::string>("in_cloud_C", "/rx_C"), pnh.param<int>("queue_size", 10)),
    sync_(SyncPolicy(pnh.param<int>("queue_size", 10)))
  {
    // 订阅话题
    in_A_ = pnh.param<std::string>("in_cloud_A", "/rx_A");
    in_C_ = pnh.param<std::string>("in_cloud_C", "/rx_C");
    out_  = pnh.param<std::string>("out_cloud", "/merged_cloud");

    // 外参（C -> A），初始值由静态参数提供
    tx_ = pnh.param<double>("tx", 0.0);
    ty_ = pnh.param<double>("ty", 0.0);
    tz_ = pnh.param<double>("tz", 0.0);
    yaw_deg_   = pnh.param<double>("yaw_deg",   0.0);
    pitch_deg_ = pnh.param<double>("pitch_deg", 0.0);
    roll_deg_  = pnh.param<double>("roll_deg",  0.0);

    // frame 名（用于 TF 与输出 header）
    frame_A_ = pnh.param<std::string>("frame_A", "camera_init_A");
    frame_C_ = pnh.param<std::string>("frame_C", "camera_init_C");

    // 同步窗口（毫秒）
    double slop_ms = pnh.param<double>("sync_slop_ms", 80.0);
    slop_sec_ = slop_ms / 1000.0;
    use_latest_stamp_ = pnh.param<bool>("use_latest_stamp", true);

    // 预计算 T(C->A)
    {
      std::lock_guard<std::mutex> lk(extrin_mtx_);
      computeTransform_unlocked();
    }

    // 绑定同步
    sync_.connectInput(sub_A_, sub_C_);
    sync_.setInterMessageLowerBound(0, ros::Duration(0.0));
    sync_.setMaxIntervalDuration(ros::Duration(slop_sec_));
    sync_.registerCallback(boost::bind(&PcMultiMerger::cb, this, _1, _2));

    pub_ = nh.advertise<sensor_msgs::PointCloud2>(out_, 1);

    // dynamic_reconfigure server
    dr_srv_.reset(new dynamic_reconfigure::Server<pc_udp_streamer::ExtrinsicConfig>(pnh));
    dynamic_reconfigure::Server<pc_udp_streamer::ExtrinsicConfig>::CallbackType cb;
    cb = [&](pc_udp_streamer::ExtrinsicConfig &cfg, uint32_t){
      std::lock_guard<std::mutex> lk(extrin_mtx_);
      tx_ = cfg.tx; ty_ = cfg.ty; tz_ = cfg.tz;
      yaw_deg_ = cfg.yaw_deg; pitch_deg_ = cfg.pitch_deg; roll_deg_ = cfg.roll_deg;
      apply_tf_ = cfg.apply_tf;
      computeTransform_unlocked();
    };
    dr_srv_->setCallback(cb);

    ROS_INFO_STREAM("pc_multi_merger: A=" << in_A_ << ", C=" << in_C_
                    << " -> " << out_ << " | init T(C->A)=[tx="<<tx_<<", ty="<<ty_
                    <<", tz="<<tz_<<", yaw="<<yaw_deg_<<", pitch="<<pitch_deg_
                    <<", roll="<<roll_deg_<<"] slop="<<slop_ms<<"ms "
                    << "frames(A="<<frame_A_<<", C="<<frame_C_<<")");
  }

private:
  void computeTransform_unlocked(){
    const double yaw=deg2rad(yaw_deg_), pitch=deg2rad(pitch_deg_), roll=deg2rad(roll_deg_);
    Eigen::AngleAxisf Ryaw (static_cast<float>(yaw),   Eigen::Vector3f::UnitZ());
    Eigen::AngleAxisf Rpitch(static_cast<float>(pitch),Eigen::Vector3f::UnitY());
    Eigen::AngleAxisf Rroll (static_cast<float>(roll), Eigen::Vector3f::UnitX());
    Eigen::Matrix3f R = (Ryaw * Rpitch * Rroll).toRotationMatrix();

    T_.setIdentity();
    T_.block<3,3>(0,0) = R;
    T_.block<3,1>(0,3) = Eigen::Vector3f(static_cast<float>(tx_),
                                         static_cast<float>(ty_),
                                         static_cast<float>(tz_));
  }

  void maybePublishTF(const ros::Time& stamp){
    if(!apply_tf_) return;
    geometry_msgs::TransformStamped ts;
    ts.header.stamp = stamp;
    ts.header.frame_id = frame_A_;
    ts.child_frame_id  = frame_C_;

    Eigen::Matrix4f T;
    {
      std::lock_guard<std::mutex> lk(extrin_mtx_);
      T = T_;
    }
    Eigen::Matrix3f R = T.block<3,3>(0,0);
    Eigen::Quaternionf q(R);
    ts.transform.translation.x = T(0,3);
    ts.transform.translation.y = T(1,3);
    ts.transform.translation.z = T(2,3);
    ts.transform.rotation.x = q.x();
    ts.transform.rotation.y = q.y();
    ts.transform.rotation.z = q.z();
    ts.transform.rotation.w = q.w();
    tf_br_.sendTransform(ts);
  }

  void cb(const sensor_msgs::PointCloud2ConstPtr& a_msg,
          const sensor_msgs::PointCloud2ConstPtr& c_msg)
  {
    // 1) 检查 layout 并取 XYZ 偏移
    FieldOffsets fa, fc;
    if (!findXYZOffsets(*a_msg, fa) || !findXYZOffsets(*c_msg, fc)) {
      ROS_WARN_THROTTLE(1.0, "PointCloud2 missing x/y/z fields");
      return;
    }
    if (a_msg->point_step != c_msg->point_step ||
        a_msg->fields.size() != c_msg->fields.size() ||
        a_msg->is_bigendian != c_msg->is_bigendian) {
      ROS_WARN_THROTTLE(1.0, "Layout mismatch between A and C; skip this pair");
      return;
    }

    // 2) 复制 C，做 C->A 刚体变换（只改 x/y/z）
    sensor_msgs::PointCloud2 c_in_A = *c_msg; // 拷贝一份
    Eigen::Matrix4f Tcur;
    {
      std::lock_guard<std::mutex> lk(extrin_mtx_);
      Tcur = T_;
    }
    try {
      transformXYZ_in_place(c_in_A, fc, Tcur);
    } catch(...) {
      ROS_WARN_THROTTLE(1.0, "transformXYZ_in_place failed");
      return;
    }

    // 3) 合并（逐字节拼接，保留所有字段）
    sensor_msgs::PointCloud2 merged;
    try {
      merged = concatCloudsSameLayout(*a_msg, c_in_A);
    } catch(const std::exception& e) {
      ROS_WARN_THROTTLE(1.0, "concat failed: %s", e.what());
      return;
    }

    // 4) header：frame_id 用 A（此时 C 已在 A 框架），stamp 取较新者/或 A
    merged.header.frame_id = a_msg->header.frame_id.empty() ? frame_A_ : a_msg->header.frame_id;
    merged.header.stamp = use_latest_stamp_
                          ? (a_msg->header.stamp > c_msg->header.stamp ? a_msg->header.stamp
                                                                        : c_msg->header.stamp)
                          : a_msg->header.stamp;

    pub_.publish(merged);

    // 5) 可选：发布 TF，便于 RViz 直观看轴
    maybePublishTF(merged.header.stamp);

    ROS_DEBUG_STREAM_THROTTLE(1.0, "[merge] A("<<a_msg->width*a_msg->height<<") + "
                               <<"C("<<c_msg->width*c_msg->height<<") -> "
                               <<merged.width*merged.height<<" pts");
  }

  // params
  std::string in_A_, in_C_, out_;
  std::string frame_A_{"camera_init_A"}, frame_C_{"camera_init_C"};
  double tx_{0}, ty_{0}, tz_{0}, yaw_deg_{0}, pitch_deg_{0}, roll_deg_{0};
  bool use_latest_stamp_{true};
  double slop_sec_{0.08};
  bool apply_tf_{true};

  // subs & sync
  message_filters::Subscriber<sensor_msgs::PointCloud2> sub_A_, sub_C_;
  message_filters::Synchronizer<SyncPolicy> sync_;
  ros::Publisher pub_;

  // extrinsics
  std::mutex extrin_mtx_;
  Eigen::Matrix4f T_{Eigen::Matrix4f::Identity()};

  // dynamic reconfigure
  std::unique_ptr<dynamic_reconfigure::Server<pc_udp_streamer::ExtrinsicConfig>> dr_srv_;

  // TF
  tf2_ros::TransformBroadcaster tf_br_;
};

int main(int argc, char** argv){
  ros::init(argc, argv, "pc_multi_merger");
  ros::NodeHandle nh, pnh("~");
  PcMultiMerger node(nh, pnh);
  ros::spin();
  return 0;
}

